#include "stdafx.h"

//#include <winsock2.h>
//#include <ws2tcpip.h>
#include <winsock/winsock_proxy.h>

#include "HookWinsock.h"

accept_proxy_t          pTrueAccept;
closesocket_proxy_t     pTrueClosesocket;
bind_proxy_t            pTrueBind;
connect_proxy_t         pTrueConnect;
freeaddrinfo_proxy_t    pTrueFreeaddrinfo;
FreeAddrInfoW_proxy_t   pTrueFreeAddrInfoW;
getaddrinfo_proxy_t     pTrueGetaddrinfo;
GetAddrInfoW_proxy_t    pTrueGetAddrInfoW;
getpeername_proxy_t     pTrueGetpeername;
getsockname_proxy_t     pTrueGetsockname;
ioctlsocket_proxy_t     pTrueIoctlsocket;
getsockopt_proxy_t      pTrueGetsockopt;
listen_proxy_t          pTrueListen;
recv_proxy_t            pTrueRecv;
recvfrom_proxy_t        pTrueRecvfrom;
select_proxy_t          pTrueSelect;
send_proxy_t            pTrueSend;
sendto_proxy_t          pTrueSendto;
shutdown_proxy_t        pTrueShutdown;
socket_proxy_t          pTrueSocket;
setsockopt_proxy_t      pTrueSetsockopt;
WSACleanup_proxy_t      pTrueWSACleanup;
__WSAFDIsSet_proxy_t    pTrue__WSAFDIsSet;
WSAGetLastError_proxy_t pTrueWSAGetLastError;
WSASetLastError_proxy_t pTrueWSASetLastError;
WSASocket_proxy_t       pTrueWSASocket;
WSAStartup_proxy_t      pTrueWSAStartup;

/////////////////////////////////////////////////////////////////////
void HookWinsock2Dll(HMODULE &hMod)
{
	pTrueAccept=(accept_proxy_t)GetProcAddress(hMod,"accept");
    if (!Mhook_SetHook((PVOID*)&pTrueAccept,accept_proxy))
    {
        throw std::runtime_error ("Mhook accept_proxy failure");
    }
    pTrueClosesocket = (closesocket_proxy_t) GetProcAddress(hMod, "closesocket");
    if (!Mhook_SetHook((PVOID*)&pTrueClosesocket, closesocket_proxy))
    {
        throw std::runtime_error ("Mhook closesocket_proxy failure");
    }
    pTrueBind = (bind_proxy_t) GetProcAddress(hMod, "bind");
    if (!Mhook_SetHook((PVOID*)&pTrueBind, bind_proxy))
    {
		throw std::runtime_error ("Mhook bind_proxy failure");
    }
    pTrueConnect = (connect_proxy_t) GetProcAddress(hMod, "connect");
    if (!Mhook_SetHook((PVOID*)&pTrueConnect, connect_proxy))
    {
		throw std::runtime_error ("Mhook connect_proxy failure");
    }
    pTrueFreeaddrinfo = (freeaddrinfo_proxy_t) GetProcAddress(hMod, "freeaddrinfo");
    if (!Mhook_SetHook((PVOID*)&pTrueFreeaddrinfo, freeaddrinfo_proxy))
    {
		throw std::runtime_error ("Mhook freeaddrinfo_proxy failure");
    }
    pTrueFreeAddrInfoW = (FreeAddrInfoW_proxy_t) GetProcAddress(hMod, "FreeAddrInfoW");
    if (!Mhook_SetHook((PVOID*)&pTrueFreeAddrInfoW, FreeAddrInfoW_proxy))
    {
		throw std::runtime_error ("Mhook FreeAddrInfoW_proxy failure");
    }
    pTrueGetaddrinfo = (getaddrinfo_proxy_t) GetProcAddress(hMod, "getaddrinfo");
    if (!Mhook_SetHook((PVOID*)&pTrueGetaddrinfo, getaddrinfo_proxy))
    {
		throw std::runtime_error ("Mhook getaddrinfo_proxy failure");
    }
    pTrueGetAddrInfoW = (GetAddrInfoW_proxy_t) GetProcAddress(hMod, "GetAddrInfoW");
    if (!Mhook_SetHook((PVOID*)&pTrueGetAddrInfoW, GetAddrInfoW_proxy))
    {
		throw std::runtime_error ("Mhook GetAddrInfoW_proxy failure");
    }
    pTrueGetpeername = (getpeername_proxy_t) GetProcAddress(hMod, "getpeername");
    if (!Mhook_SetHook((PVOID*)&pTrueGetpeername, getpeername_proxy))
    {
        throw std::runtime_error ("Mhook getpeername_proxy failure");
    }
    pTrueGetsockname = (getsockname_proxy_t) GetProcAddress(hMod, "getsockname");
    if (!Mhook_SetHook((PVOID*)&pTrueGetsockname, getsockname_proxy))
    {
		throw std::runtime_error ("Mhook getsockname_proxy failure");
    }
    pTrueIoctlsocket = (ioctlsocket_proxy_t) GetProcAddress(hMod, "ioctlsocket");
    if (!Mhook_SetHook((PVOID*)&pTrueIoctlsocket, ioctlsocket_proxy))
    {
		throw std::runtime_error ("Mhook ioctlsocket_proxy failure");
    }
    pTrueGetsockopt = (getsockopt_proxy_t) GetProcAddress(hMod, "getsockopt");
    if (!Mhook_SetHook((PVOID*)&pTrueGetsockopt, getsockopt_proxy))
    {
		throw std::runtime_error ("Mhook getsockopt_proxy failure");
    }
    pTrueListen = (listen_proxy_t) GetProcAddress(hMod, "listen");
    if (!Mhook_SetHook((PVOID*)&pTrueListen, bind_proxy))
    {
		throw std::runtime_error ("Mhook listen_proxy failure");
    }
    pTrueRecv = (recv_proxy_t) GetProcAddress(hMod, "recv");
    if (!Mhook_SetHook((PVOID*)&pTrueRecv, recv_proxy))
    {
		throw std::runtime_error ("Mhook recv_proxy failure");
    }
    pTrueRecvfrom = (recvfrom_proxy_t) GetProcAddress(hMod, "recvfrom");
    if (!Mhook_SetHook((PVOID*)&pTrueRecvfrom, recvfrom_proxy))
    {
		throw std::runtime_error ("Mhook recvfrom_proxy failure");
    }
    pTrueSelect = (select_proxy_t) GetProcAddress(hMod, "select");
    if (!Mhook_SetHook((PVOID*)&pTrueSelect, select_proxy))
    {
		throw std::runtime_error ("Mhook select_proxy failure");
    }
    pTrueSend = (send_proxy_t) GetProcAddress(hMod, "send");
    if (!Mhook_SetHook((PVOID*)&pTrueSend, send_proxy))
    {
		throw std::runtime_error ("Mhook send_proxy failure");
    }
    pTrueSendto = (sendto_proxy_t) GetProcAddress(hMod, "sendto");
    if (!Mhook_SetHook((PVOID*)&pTrueSendto, sendto_proxy))
    {
		throw std::runtime_error ("Mhook sendto_proxy failure");
    }
    pTrueShutdown = (shutdown_proxy_t) GetProcAddress(hMod, "shutdown");
    if (!Mhook_SetHook((PVOID*)&pTrueShutdown, shutdown_proxy))
    {
		throw std::runtime_error ("Mhook shutdown_proxy failure");
    }
    pTrueSocket = (socket_proxy_t) GetProcAddress(hMod, "socket");
    if (!Mhook_SetHook((PVOID*)&pTrueSocket, socket_proxy))
    {
		throw std::runtime_error ("Mhook socket_proxy failure");
    }
    pTrueSetsockopt = (setsockopt_proxy_t) GetProcAddress(hMod, "setsockopt");
    if (!Mhook_SetHook((PVOID*)&pTrueSetsockopt, setsockopt_proxy))
    {
		throw std::runtime_error ("Mhook setsockopt_proxy failure");
    }
    pTrueWSACleanup = (WSACleanup_proxy_t) GetProcAddress(hMod, "WSACleanup");
    if (!Mhook_SetHook((PVOID*)&pTrueWSACleanup, WSACleanup_proxy))
    {
		throw std::runtime_error ("Mhook WSACleanup_proxy failure");
    }
    pTrue__WSAFDIsSet = (__WSAFDIsSet_proxy_t) GetProcAddress(hMod, "__WSAFDIsSet");
    if (!Mhook_SetHook((PVOID*)&pTrue__WSAFDIsSet, __WSAFDIsSet_proxy))
    {
		throw std::runtime_error ("Mhook __WSAFDIsSet_proxy failure");
    }
    pTrueWSAGetLastError = 
        (WSAGetLastError_proxy_t) GetProcAddress(hMod, "WSAGetLastError");
    if (!Mhook_SetHook((PVOID*)&pTrueWSAGetLastError, WSAGetLastError_proxy))
    {
		throw std::runtime_error ("Mhook WSAGetLastError_proxy failure");
    }
    pTrueWSASetLastError = 
        (WSASetLastError_proxy_t)GetProcAddress(hMod,"WSASetLastError");
	if (!Mhook_SetHook((PVOID*)&pTrueWSASetLastError,WSASetLastError_proxy))
    {
		throw std::runtime_error ("Mhook WSASetLastError_proxy failure");
    }
    pTrueWSASocket = (WSASocket_proxy_t) GetProcAddress(hMod, "WSASocket");
    if (!Mhook_SetHook((PVOID*)&pTrueWSASocket, WSASocket_proxy))
    {
		throw std::runtime_error ("Mhook WSASocket_proxy failure");
    }
    pTrueWSAStartup = (WSAStartup_proxy_t) GetProcAddress(hMod, "WSAStartup");
    if (!Mhook_SetHook((PVOID*)&pTrueWSAStartup, WSAStartup_proxy))
    {
		throw std::runtime_error ("Mhook WSAStartup_proxy failure");
    }

    ConnectToVirtualChannel();
}

/////////////////////////////////////////////////////////////////////
void UnhookWinsock2Dll ()
{
	Mhook_Unhook((PVOID*)&pTrueAccept);
	Mhook_Unhook((PVOID*)&pTrueClosesocket);
	Mhook_Unhook((PVOID*)&pTrueBind);
	Mhook_Unhook((PVOID*)&pTrueConnect);
    Mhook_Unhook((PVOID*)&pTrueFreeaddrinfo);
    Mhook_Unhook((PVOID*)&pTrueFreeAddrInfoW);
    Mhook_Unhook((PVOID*)&pTrueGetaddrinfo);
    Mhook_Unhook((PVOID*)&pTrueGetAddrInfoW);
	Mhook_Unhook((PVOID*)&pTrueGetpeername);
	Mhook_Unhook((PVOID*)&pTrueGetsockname);
	Mhook_Unhook((PVOID*)&pTrueIoctlsocket);
	Mhook_Unhook((PVOID*)&pTrueGetsockopt);
	Mhook_Unhook((PVOID*)&pTrueListen);
	Mhook_Unhook((PVOID*)&pTrueRecv);
	Mhook_Unhook((PVOID*)&pTrueRecvfrom);
	Mhook_Unhook((PVOID*)&pTrueSelect);
	Mhook_Unhook((PVOID*)&pTrueSend);
	Mhook_Unhook((PVOID*)&pTrueSendto);
	Mhook_Unhook((PVOID*)&pTrueShutdown);
	Mhook_Unhook((PVOID*)&pTrueSocket);
	Mhook_Unhook((PVOID*)&pTrueSetsockopt);
	Mhook_Unhook((PVOID*)&pTrueWSACleanup);
	Mhook_Unhook((PVOID*)&pTrue__WSAFDIsSet);
	Mhook_Unhook((PVOID*)&pTrueWSAGetLastError);
	Mhook_Unhook((PVOID*)&pTrueWSASetLastError);
    Mhook_Unhook((PVOID*)&pTrueWSASocket);
	Mhook_Unhook((PVOID*)&pTrueWSAStartup);
}
